from common.variables import TimeFrameType
import numpy as np
import os
import pandas as pd
import plotly.graph_objects as go
import random
from trade.strategy import BaseStrategy
from trade.tools import SlidingData


class RandomStrategy(BaseStrategy):
    need_symbols: list[str]
    preloads: int = 20
    timeframe: TimeFrameType = "1m"

    def __init__(self, use_symbol: str = "BTC/USDT:USDT"):
        self.use_symbol = use_symbol
        self.need_symbols = [use_symbol, "BTC/USDT"]

    def apply(self):
        closes = self.broker.data_records[(self.use_symbol, "close")][-10:]
        opens = self.broker.data_records[(self.use_symbol, "open")][-10:]
        yangs = [closes[i] > opens[i] for i in range(len(closes))]
        positions = self.broker.get_positions()
        yang_cnt = sum(yangs)
        if yang_cnt > 5:
            if self.use_symbol not in positions:
                size = 1
            elif positions[self.use_symbol]["side"] == "long":
                size = positions[self.use_symbol]["contracts"] + 1
            elif (
                positions[self.use_symbol]["side"] == "short"
                and positions[self.use_symbol]["contracts"] == 0
            ):
                size = 1
            else:
                size = 0
            if size > 0:
                print(f"下单：{self.use_symbol}, sell, {size}")
                self.broker.order_future(
                    symbol=self.use_symbol,
                    size=size,
                    price=None,
                    timeInForce="GTC",
                    market_type="market",
                )

        if yang_cnt < 5:
            if self.use_symbol not in positions:
                size = 1
            elif positions[self.use_symbol]["side"] == "short":
                size = positions[self.use_symbol]["contracts"] + 1
            elif (
                positions[self.use_symbol]["side"] == "long"
                and positions[self.use_symbol]["contracts"] == 0
            ):
                size = 1
            else:
                size = 0
            if size > 0:
                print(f"下单：{self.use_symbol}, buy, {size}")
                self.broker.order_future(
                    symbol=self.use_symbol,
                    size=size,
                    price=None,
                    timeInForce="GTC",
                    market_type="market",
                )


def limit_buy(x):
    if x < 5.5:
        return 0
    elif x > 10:
        return 10
    return x


class RandomBuyStrategy(BaseStrategy):
    def __init__(
        self, use_spot: str = "FDUSD/USDT", use_future: str = "USDC/USDT:USDT"
    ):
        self.use_spot = use_spot
        self.use_future = use_future

    def apply(self):
        unimrr = self.broker.get_unimrr()
        x = random.random()
        if unimrr > 10 and x > 0.5:
            free_margin = self.broker.get_free_margin("USDT")
            buy_spot = (
                free_margin
                * 0.4
                * random.random()
                / self.broker.data_records[(self.use_spot, "close")][-1]
                * self.broker.get_leverage(self.use_spot)
            )
            add_future = (
                free_margin
                * 0.3
                * random.random()
                / self.broker.data_records[(self.use_future, "close")][-1]
                * self.broker.get_leverage(self.use_future)
            )
            buy_spot = limit_buy(round(buy_spot, 0))
            add_future = limit_buy(round(add_future, 0))
            if buy_spot > 0:
                self.broker.order_spot(self.use_spot, buy_spot, None, "GTC", "market")
            if add_future > 0:
                try:
                    if (
                        self.use_future in self.broker.positions
                        and self.broker.positions[self.use_future]["side"] == "long"
                    ):
                        self.broker.order_future(
                            self.use_future, add_future, None, "GTC", "market"
                        )
                    else:
                        self.broker.order_future(
                            self.use_future, add_future, None, "GTC", "market"
                        )
                except Exception as e:
                    print(f"debug {e}")
        else:
            balance = self.use_spot.split("/")[0]
            if balance in self.broker.balances:
                sell_spot = (
                    self.broker.balances[balance]["total"] * random.random() * 0.5
                )
                sell_spot = limit_buy(round(sell_spot, 0))
                if sell_spot > 0:
                    self.broker.order_spot(
                        self.use_spot, sell_spot, None, "GTC", "market"
                    )
            if self.use_future in self.broker.positions:
                sell_future = (
                    self.broker.positions[self.use_future]["contracts"]
                    * random.random()
                    * 0.5
                )
                sell_future = limit_buy(round(sell_future, 0))
                if sell_future > 0:
                    try:
                        if (
                            self.use_future in self.broker.positions
                            and self.broker.positions[self.use_future]["side"] == "long"
                        ):
                            self.broker.order_future(
                                self.use_future,
                                sell_future,
                                None,
                                "GTC",
                                "market",
                            )
                        else:
                            self.broker.order_future(
                                self.use_future,
                                sell_future,
                                None,
                                "GTC",
                                "market",
                            )
                    except Exception as e:
                        print(f"debug {e}")


class RandomStrategySpot(BaseStrategy):
    def __init__(self, use_symbol: str = "BTC/USDT"):
        self.use_symbol = use_symbol

    def apply(self):
        use_balance = self.broker.markets[self.use_symbol]["base"]
        closes = self.broker.get_data(self.use_symbol, "close", slice(-10, None))
        opens = self.broker.get_data(self.use_symbol, "open", slice(-10, None))
        yangs = [closes[i] > opens[i] for i in range(len(closes))]
        balances = self.broker.get_balances()
        yang_cnt = sum(yangs)
        target = 5
        if yang_cnt > 6:
            if use_balance not in balances:
                size = target
            elif balances[use_balance]["total"] < target:
                size = target - balances[use_balance]["total"]
            else:
                size = 0

            if size > 0:
                self.broker.order_spot(
                    self.use_symbol,
                    size=size,
                    price=None,
                    timeInForce="GTC",
                    market_type="market",
                )

        if yang_cnt < 4:
            if use_balance not in balances:
                size = 0
            elif balances[use_balance]["total"] > 0:
                size = balances[use_balance]["total"]
            else:
                size = 0

            if size > 0:
                self.broker.order_spot(
                    self.use_symbol,
                    size=size,
                    price=None,
                    timeInForce="GTC",
                    market_type="market",
                )


class TwoLineStrategy(BaseStrategy):
    def __init__(self, use_spot: str = "BTC/USDT", model_file: str = None):
        self.use_spot = use_spot
        self.records = []
        self.slide_data = SlidingData(200, [7, 9, 25, 90])
        self.model = None

    def apply(self):
        unimrr = self.broker.get_unimrr()
        cur_close = self.broker.get_data(self.use_spot, "close", -1)
        self.slide_data.push(cur_close)
        ma7 = self.slide_data.get_value(7, slice(-4)).__reversed__()
        ma9 = self.slide_data.get_value(9, slice(-4)).__reversed__()
        ma25 = self.slide_data.get_value(25, slice(-4)).__reversed__()
        ma90 = self.slide_data.get_value(90, slice(-4)).__reversed__()
        closes = self.broker.get_data(
            self.use_spot, "close", slice(-16, -1)
        ).__reversed__()
        features = [
            self.broker.get_data(self.use_spot, "open", -1),
            self.broker.get_data(self.use_spot, "close", -1),
            self.broker.get_data(self.use_spot, "high", -1),
            self.broker.get_data(self.use_spot, "low", -1),
            *ma7,
            *ma9,
            *ma25,
            *ma90,
            *closes,
        ]
        deal_feature = np.array(features) / cur_close
        result = self.model.predict(deal_feature[np.newaxis])
        neg1 = features[4] - features[12]
        neg2 = features[5] - features[13]
        xtime = self.broker.get_cur_time()
        print(f"{xtime} signal {neg1} {neg2}")
        dr = 0
        if unimrr > 10:
            free_margin = self.broker.get_free_margin(
                self.broker.markets[self.use_spot]["quote"]
            )
            size = (
                free_margin
                * self.broker.get_leverage(self.use_spot)
                * 0.8
                / self.broker.get_data(self.use_spot, "close", 0)
            )
            if neg1 > 0 and neg2 < 0 and result > 1e-5:
                self.broker.adjust_to_target_spot(
                    self.use_spot, size, None, "GTC", "market"
                )
                dr = 1
            elif neg1 < 0 and neg2 > 0 and result < -1e-5:
                self.broker.adjust_to_target_spot(
                    self.use_spot, -size, None, "GTC", "market"
                )
                dr = -1
        else:
            self.broker.close_position(
                self.use_spot, None, timeInForce="GTC", market_type="market"
            )

        self.records.append(
            [
                xtime,
                self.broker.get_data(self.use_spot, "close", 0),
                self.broker.get_account_value(),
                dr,
            ]
        )

    def after_all(self):
        basedir = "./rub"
        os.makedirs(basedir, exist_ok=True)
        df = pd.DataFrame(data=self.records, columns=["time", "BTC", "value", "dr"])
        df.to_csv(os.path.join(basedir, "aa.csv"))
        fig = go.Figure()

        fig.add_trace(
            go.Scatter(
                x=df["time"],
                y=df["BTC"] / df["BTC"][0],
                name="比特币",
                yaxis="y1",
                mode="lines",
            )
        )

        fig.add_trace(
            go.Scatter(
                x=df["time"],
                y=df["value"] / df["value"][0],
                name="策略",
                yaxis="y1",
                mode="lines",
            )
        )

        fig.add_trace(go.Bar(x=df["time"], y=df["dr"], name="方向信号", yaxis="y2"))

        fig.update_layout(
            title="策略和基准对比图",
            xaxis_title="日期",
            yaxis_title="净值",
            xaxis_type="category",
            xaxis_rangeslider_visible=True,
            height=1000,
            yaxis2=dict(
                title="增长率",
                overlaying="y",
                side="right",
                fixedrange=False,
                zeroline=True,  # 显示x轴的0轴线
                zerolinewidth=2,  # 设置0轴线宽度
                zerolinecolor="gray",
            ),
            yaxis=dict(fixedrange=False),
        )
        fig.write_html(os.path.join(basedir, "result.html"))
